package com.example.server.controller;

import com.example.server.entity.ChatMessage;
import com.example.server.util.SignatureUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/chat")
public class ChatController {

    private final WebClient webClient;
    private final SignatureUtil signatureUtil;

    private final Map<String, Boolean> sessionInitialPromptSent = new ConcurrentHashMap<>();

    @Autowired
    public ChatController(WebClient webClient, SignatureUtil signatureUtil) {
        this.webClient = webClient;
        this.signatureUtil = signatureUtil;
    }

    @PostMapping("/send")
    public Mono<ResponseEntity<String>> sendMessage(@RequestBody ChatMessage message) {
        String sessionId = message.getSessionId();
        String prompt = message.getContent();

        // Check if this is the first message for the session
        if (!sessionInitialPromptSent.containsKey(sessionId)) {
            prompt = "假设我们俩都是小说家，我们现在就写小说相关的问题进行讨论，我希望你可以就我告诉你的内容给我一些写小说方面的灵感。" + prompt;
            sessionInitialPromptSent.put(sessionId, true);
        }

        Map<String, String> queryParams = new HashMap<>();
        queryParams.put("requestId", UUID.randomUUID().toString());

        MultiValueMap<String, String> multiValueMap = new LinkedMultiValueMap<>();
        queryParams.forEach((key, value) -> multiValueMap.add(key, value));

        System.out.println("Received message from frontend:");
        System.out.println("Content: " + message.getContent());

        return webClient.post()
                .uri(uriBuilder -> uriBuilder.path("/vivogpt/completions")
                        .queryParams(multiValueMap)
                        .build())
                .headers(headers -> headers.setAll(signatureUtil.generateAuthHeaders("POST", "/vivogpt/completions", queryParams)))
                .bodyValue(buildRequestBody(sessionId, prompt))
                .retrieve()
                .bodyToMono(String.class)
                .map(body -> {
                    System.out.println("AIGC Feedback:");
                    System.out.println(body); // Print the feedback
                    return ResponseEntity.ok().body(body);
                })
                .onErrorResume(e -> {
                    e.printStackTrace();
                    return Mono.just(ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR).body("Error processing your request: " + e.getMessage()));
                });
    }

    private String buildRequestBody(String sessionId, String prompt) {
        return String.format("{\"prompt\": \"%s\", \"model\": \"vivo-BlueLM-TB\", \"sessionId\": \"%s\"}", prompt, sessionId);
    }
}
